package com.ics.atable.chat.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.ics.atable.chat.expection.ErrorCode;
import com.ics.atable.chat.expection.ThrowUtils;
import com.ics.atable.chat.mapper.EmbeddingsPersistMapper;
import com.ics.atable.chat.model.entity.EmbeddingsPersist;
import com.ics.atable.chat.service.EmbeddingsPersistService;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.ObjectUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Service
public class EmbeddingsPersistServiceImpl extends ServiceImpl<EmbeddingsPersistMapper, EmbeddingsPersist>
    implements EmbeddingsPersistService {

    @Autowired
    private EmbeddingsPersistMapper embeddingsPersistMapper;

    @Override
    public float[] getEmbeddingVectorByContent(String content, String type) {
        ThrowUtils.throwIf(StringUtils.isEmpty(content), ErrorCode.PARAMS_ERROR, "内容不能为空");
        ThrowUtils.throwIf(StringUtils.isEmpty(type), ErrorCode.PARAMS_ERROR, "类型不能为空");

        LambdaQueryWrapper<EmbeddingsPersist> queryWrapper = new LambdaQueryWrapper<>();
        queryWrapper.eq(EmbeddingsPersist::getContent, content)
                .eq(EmbeddingsPersist::getType, type);
        EmbeddingsPersist embeddingsPersist = embeddingsPersistMapper.selectOne(queryWrapper);
        if (embeddingsPersist == null) {
            return null;
        }
        return embeddingsPersist.getVector();
    }

    @Override
    public boolean saveOrUpdateEmbeddingVector(EmbeddingsPersist embeddingsPersist) {
        ThrowUtils.throwIf(ObjectUtils.isEmpty(embeddingsPersist), ErrorCode.PARAMS_ERROR, "参数不能为空");
        ThrowUtils.throwIf(StringUtils.isEmpty(embeddingsPersist.getContent()), ErrorCode.PARAMS_ERROR, "内容不能为空");
        ThrowUtils.throwIf(StringUtils.isEmpty(embeddingsPersist.getType()), ErrorCode.PARAMS_ERROR, "类型不能为空");
        ThrowUtils.throwIf(ObjectUtils.isEmpty(embeddingsPersist.getVector()), ErrorCode.PARAMS_ERROR, "向量不能为空");
        ThrowUtils.throwIf(ObjectUtils.isEmpty(embeddingsPersist.getDimensions()), ErrorCode.PARAMS_ERROR, "向量维度不能为空");

        // 根据content字段查询是否已存在记录
        EmbeddingsPersist existingRecord = this.lambdaQuery()
                .eq(EmbeddingsPersist::getContent, embeddingsPersist.getContent())
                .eq(EmbeddingsPersist::getType, embeddingsPersist.getType())
                .one();
        if (existingRecord != null) {
            // 如果存在，则更新向量数据
            embeddingsPersist.setId(existingRecord.getId());
            return this.updateById(embeddingsPersist);
        } else {
            // 如果不存在，则插入新记录
            return this.save(embeddingsPersist);
        }
    }


    @Override
    public Map<String, float[]> getEmbeddingVectorsByContents(List<String> contents, String type) {
        ThrowUtils.throwIf(CollectionUtils.isEmpty(contents), ErrorCode.PARAMS_ERROR, "内容列表不能为空");
        ThrowUtils.throwIf(StringUtils.isEmpty(type), ErrorCode.PARAMS_ERROR, "类型不能为空");

        // 去重并查询
        List<String> uniqueContents = contents.stream().distinct().collect(Collectors.toList());

        List<EmbeddingsPersist> embeddingsPersists = this.lambdaQuery()
                .in(EmbeddingsPersist::getContent, uniqueContents)
                .eq(EmbeddingsPersist::getType, type)
                .list();

        // 按原始顺序返回结果，未找到的返回null
        return embeddingsPersists.stream()
                .collect(Collectors.toMap(EmbeddingsPersist::getContent, EmbeddingsPersist::getVector));
    }

    @Override
    @Transactional(rollbackFor = Exception.class)
    public void saveOrUpdateBatchEmbeddingVectors(List<EmbeddingsPersist> embeddingsPersists) {
        ThrowUtils.throwIf(CollectionUtils.isEmpty(embeddingsPersists), ErrorCode.PARAMS_ERROR, "向量列表不能为空");

        // 获取所有content
        List<String> contents = embeddingsPersists.stream()
                .map(EmbeddingsPersist::getContent)
                .distinct()
                .collect(Collectors.toList());

        // 查询已存在的记录
        List<EmbeddingsPersist> existingRecords = this.lambdaQuery()
                .in(EmbeddingsPersist::getContent, contents)
                .eq(EmbeddingsPersist::getType, embeddingsPersists.getFirst().getType())
                .list();

        // 创建content到id的映射
        Map<String, Long> contentToIdMap = existingRecords.stream()
                .collect(Collectors.toMap(EmbeddingsPersist::getContent, EmbeddingsPersist::getId));

        // 设置已存在记录的ID
        for (EmbeddingsPersist embeddingsPersist : embeddingsPersists) {
            Long id = contentToIdMap.get(embeddingsPersist.getContent());
            if (id != null) {
                embeddingsPersist.setId(id);
            }
        }

        // 批量保存或更新
        this.saveOrUpdateBatch(embeddingsPersists);
    }


}